package a4ny;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.Set;
import java.util.TreeMap;

public class ViterbiTaggerPerplexity {
	private TreeMap<String, TreeMap<String, Double>> wordPerp;
	private TreeMap<String, Double> posPerp;
	private TreeMap<Bigram, Double> bigramPerp;
	
	public ViterbiTaggerPerplexity(TreeMap<String, TreeMap<String, Double> > wordPerp, 
			TreeMap<String, Double> posPerp, TreeMap<Bigram, Double> bigramPerp){
		this.wordPerp = wordPerp;
		this.posPerp = posPerp;
		this.bigramPerp = bigramPerp;
	}
	
	public ArrayList<String> tag(ArrayList<String> words){
		// Make an array of the pos tags
		Set<String> posUnique = posPerp.keySet();
		Object[] posTemp = posUnique.toArray();
		String[] pos = new String[posTemp.length];
		for(int i = 0; i < pos.length; ++i){
			pos[i] = (String)posTemp[i];
		}
		
		// Initialize the table
		double[][] table = new double[pos.length][words.size()];
		for(int i = 0; i < pos.length; ++i){
			for(int j = 0; j < words.size(); ++j){
				// Arbitrarily chosen placeholder corresponding to 0 for probabilities
				table[i][j] = 20.0;
			}
		}
		
		// Run the Viterbi algorithm for each column of the table
		int[][] backPointers = new int[pos.length][words.size()];
		for(int i = 0; i < words.size(); ++i){
			viterbi(words.get(i), i, table, backPointers, pos);
		}
		
		// Identify the most likely element of the last column,
		// then find the pos tag sequence by following the back pointers
		LinkedList<String> ppos = new LinkedList<String>();
		double minPerp = 20.0;
		int minPerpIndex = 0;
		for(int col = words.size() - 1; col >= 0; --col){
			if(col == words.size() - 1){
				for(int row = 0; row < pos.length; ++row){
					double thisPerp = table[row][col];
					if(thisPerp < minPerp){
						minPerp = thisPerp;
						minPerpIndex = row;
					}
				}
			}
			ppos.addFirst(pos[minPerpIndex]);
			minPerpIndex = backPointers[minPerpIndex][col];
		}
		return new ArrayList<String>(ppos);
	}
	
	private double getBigramPerp(String variable, String condition){
		Bigram b = new Bigram(variable, condition);
		if(bigramPerp.containsKey(b)){
			return bigramPerp.get(b);
		}else if(posPerp.containsKey(variable)){
			// Back off to unigrams
			return posPerp.get(variable);
		}else{
			return 20.0;
		}
	}
	
	private void viterbi(String word, int wordIndex, double[][] table, int[][] backPointers, String[] pos){
		// If we have no record of the word we replace the perplexities Perp(word|tags) with the overall perplexities of the pos instead
		TreeMap<String, Double> thisWordPerp;
		if(wordPerp.containsKey(word)){
			thisWordPerp = wordPerp.get(word);
		}else{
			thisWordPerp = posPerp;
		}
		
		// Initialize the first column of the matrix
		if(wordIndex == 0){			
			for(int j = 0; j < pos.length; ++j){
				if(thisWordPerp.containsKey(pos[j])){
					table[j][0] = thisWordPerp.get(pos[j])*getBigramPerp(pos[j], "<s>");
				}
			}
		// Computations for the rest of the columns
		}else{
			for(int j = 0; j < pos.length; ++j){
				double minPerp = 20.0;
				int minPerpIndex = -1;
				for(int i = 0; i < pos.length; ++i){
					if(table[i][wordIndex-1] == 20.0) continue;
					double thisWordPosPerp = thisWordPerp.containsKey(pos[j]) ? thisWordPerp.get(pos[j]) : 20.0;
					double thisPerp = getBigramPerp(pos[j], pos[i])*thisWordPosPerp*table[i][wordIndex-1];
					if(thisPerp < minPerp){
						minPerp = thisPerp;
						minPerpIndex = i;
					}
				}
				table[j][wordIndex] = minPerp;
				backPointers[j][wordIndex] = minPerpIndex;
			}
		}
	}

}
